import asyncio
import re

from app.api.routes.websocket import build_prompt, websocket_chat, llm_streamer
from app.vectorDB.chromaDB import vectorDB


# 提示词模板
prompt_template = """
你是一个辩手。
你的任务是根据下面给定的已知信息与用户进行辩论。

已知信息：
{context}  

用户发言：
{query}  

辩论时请使用中文。
"""


class RAG_BOT:
    def __init__(self, vector_db, api, n=2):
        self.vector_db = vector_db
        self.api = api
        self.n = n

    async def chat(self, query):
        # 检索
        search_results = self.vector_db.search(query, self.n)
        # 构建提示词
        prompt = build_prompt(
            prompt_template, context=search_results['documents'][0], query=query
        )
        # print(prompt)

        # 构建符合 llm_streamer 要求的消息数组
        messages = [{"role": "user", "content": prompt}]

        # 调用 deepseek api
        response = ""
        async for chunk in self.api(messages):  # 正确迭代异步生成器
            response += chunk
        return response


# 创建一个 RAG ROBOT
robot = RAG_BOT(
    vectorDB,   # 向量数据库实例
    api=llm_streamer  # 大模型调用方法
)

async def main():
    query = "爱需要克制吗？"

    response = await robot.chat(query)

    # print(response)

    cleaned_response = re.sub(r'<async_generator object.*?>$', '', str(response))
    print(cleaned_response)

asyncio.run(main())









